import json
import traceback

from flask import Flask, request

import services
from message import *

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024


@app.route("/model/train/start", methods=["POST"])
def train_start():
    """
        modelName:string
        dataset: dataset.value, 				// string
        epochs: epochs.value, 				// number
        trainRatio: trainRatio.value, 		// number
        testRatio: testRatio.value, 			// number
        valRatio: valRatio.value, 			// number
        batchSize: batchSize.value, 			// number
        learningRate: learningRate.value, 	// number
        optimizer: optimizer.value, 			// string
    :return:
    """
    train_params = request.json
    print(train_params)
    try:
        services.train_model(
            model_name=train_params["modelName"],
            dataset_id=train_params["dataset"],
            num_epochs=train_params["epochs"],
            batch_size=train_params["batchSize"],
            learning_rate=train_params["learningRate"],
            optimizer=train_params["optimizer"].lower(),
            train_ratio=train_params["trainRatio"],
            val_ratio=train_params["valRatio"],
            test_ratio=train_params["testRatio"]
        )
        return response.ok()
    except ValueError:
        traceback.print_exc()
        return response.error("Error starting train.")


@app.route("/model/train/progress/<model_name>", methods=["GET"])
def progress(model_name):
    try:
        progresses, finished = services.get_train_process(model_name=model_name)
        _time, _loss, _train_accuracy, _val_accuracy = [], [], [], []
        for item in progresses:
            _time.append(item.now_epoch)
            _loss.append(item.train_loss)
            _train_accuracy.append(item.train_acc)
            _val_accuracy.append(item.val_acc)
        ret = {
            "time": _time,
            "loss": _loss,
            "trainAccuracy": _train_accuracy,
            "valAccuracy": _val_accuracy,
            "finished": finished
        }
        print(json.dumps(ret))
        return response.ok(json.dumps(ret))
    except ValueError:
        traceback.print_exc()
        return response.error(f"Error getting progress.")


@app.route("/model/test", methods=["POST"])
def test_img():
    try:
        test_params = request.form.to_dict()
        model_name = test_params["modelName"]
        img_file = request.files['image']
        return str(services.test_img(model_name, img_file))
    except Exception:
        traceback.print_exc()
        return response.error(f"Error testing image.")


if __name__ == '__main__':
    app.run()
